// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/init/init.h>

#include <pollux/exec/plan_node_stats.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/parse/expressions.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::plan;
using namespace kumo::pollux::exec::test;

class ConjunctTest : public OperatorTestBase {
};

TEST_F(ConjunctTest, mixed) {
    auto mixed = make_row_vector(
        {
            make_flat_vector<int32_t>(18, [](auto row) { return row; }),
            make_flat_vector<bool>(
                18,
                [](auto row) { return row % 3 == 0; },
                [](auto row) { return row % 3 == 2; }),
            make_flat_vector<bool>(
                18,
                [](auto row) { return (row / 3) % 3 == 0; },
                [](auto row) { return (row / 3) % 3 == 2; })
        });
    auto plan =
            PlanBuilder()
            .values({mixed})
            .project({"c0", "c1", "c2", "if (c0 <9, c1 and c2, c1 or c2)"})
            .planNode();
    createDuckDbTable({mixed});

    assertQuery(
        plan, "select c0, c1, c2, if (c0 < 9, c1 and c2, c1 or c2) from tmp");

    // Non-null and nullable
    auto mixedOneNull = make_row_vector(
        {
            make_flat_vector<int32_t>(18, [](auto row) { return row; }),
            make_flat_vector<bool>(18, [](auto row) { return row % 3 == 0; }),
            make_flat_vector<bool>(
                18,
                [](auto row) { return (row / 3) % 3 == 0; },
                [](auto row) { return (row / 3) % 3 == 2; })
        });
    plan = PlanBuilder()
            .values({mixedOneNull})
            .project({"c0", "c1", "c2", "if (c0 <9, c1 and c2, c1 or c2)"})
            .planNode();
    createDuckDbTable({mixedOneNull});

    assertQuery(
        plan, "select c0, c1, c2, if (c0 < 9, c1 and c2, c1 or c2) from tmp");

    // Both are non-null
    auto mixedNonNull = make_row_vector(
        {
            make_flat_vector<int32_t>(18, [](auto row) { return row; }),
            make_flat_vector<bool>(18, [](auto row) { return row % 3 == 0; }),
            make_flat_vector<bool>(18, [](auto row) { return (row / 3) % 3 == 0; })
        });
    plan = PlanBuilder()
            .values({mixedNonNull})
            .project({"c0", "c1", "c2", "if (c0 <9, c1 and c2, c1 or c2)"})
            .planNode();
    createDuckDbTable({mixedNonNull});

    assertQuery(
        plan, "select c0, c1, c2, if (c0 < 9, c1 and c2, c1 or c2) from tmp");
}

TEST_F(ConjunctTest, constant) {
    for (auto counter = 0; counter < 9; ++counter) {
        auto allSame = make_row_vector(
            {
                make_flat_vector<int32_t>(18, [&](auto row) { return row; }),
                make_flat_vector<bool>(
                    18,
                    [&](auto row) { return counter % 3 == 0; },
                    [&](auto row) { return counter % 3 == 2; }),
                make_flat_vector<bool>(
                    18,
                    [&](auto row) { return (counter / 3) % 3 == 0; },
                    [&](auto row) { return (counter / 3) % 3 == 2; })
            });
        auto plan =
                PlanBuilder()
                .values({allSame})
                .project({"c0", "c1", "c2", "if (c0 <9, c1 and c2, c1 or c2)"})
                .planNode();
        createDuckDbTable({allSame});

        assertQuery(
            plan, "select c0, c1, c2, if (c0 < 9, c1 and c2, c1 or c2) from tmp");
    }
}
